package cn.turboinfo.dongying.api.provider.admin.component.security;

import cn.turboinfo.dongying.api.entity.admin.exception.security.RateLimitReachedException;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.Lists;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

@Slf4j
@RequiredArgsConstructor
public class RateLimitHelper {

    private final LoadingCache<String, List<Long>> rateLimiterCache = CacheBuilder.newBuilder()
            .maximumSize(1000)
            .build(new CacheLoader<String, List<Long>>() {
                @Override
                public List<Long> load(@SuppressWarnings("NullableProblems") String username) throws Exception {
                    return Lists.newArrayList(-1L, -1L);
                }
            });

    public void checkLimit(String key, long timeout, TimeUnit timeUnit, long retryTimes) throws RateLimitReachedException {
        try {
            List<Long> limitData = rateLimiterCache.get(key);

            long lastRefillTime = limitData.get(0);
            long remainingTokens = limitData.get(1);

            long refillTime = System.currentTimeMillis();

            if (lastRefillTime < 0 || remainingTokens < 0) {
                // 初始化填满桶
                remainingTokens = retryTimes;
            } else {
                long intervalSinceLast = refillTime - lastRefillTime;
                if (intervalSinceLast > timeUnit.toMillis(timeout)) {
                    // 距离上次填充已经超过了限制期, 直接填满
                    remainingTokens = retryTimes;
                } else {
                    // 按照时间平均分布计算应该填充的 Token 数量
                    long grantedTokens = (long) (intervalSinceLast / ((double) timeUnit.toMillis(timeout) / retryTimes));
                    remainingTokens = Math.min(grantedTokens + remainingTokens, retryTimes);
                }
            }

            // 如果超出了限制就说明本次已经超限了, 抛出异常
            if (remainingTokens <= 0) {
                throw new RateLimitReachedException();
            }

            // 最后更新值和有效期
            rateLimiterCache.put(key, Lists.newArrayList(refillTime, remainingTokens - 1));
        } catch (ExecutionException e) {
            log.error(e.getMessage(), e);
        }
    }

}
